using System;
using System.Collections.Generic;
using System.Globalization;
using System.IO;

using Atomic.Libraries;
using Atomic.Libraries.Mathematics;
using Atomic.Structures;

namespace Atomic.Vasp
{
	public static class VaspPoscar
	{
		/// <summary>
		/// Writes VASPs POSCAR file. Proper sorting of the structure is handled automatically.
		/// </summary>
		public static void Write(Structure structure, FileInfo file)
		{
			IEnumerable<Site> sortedSites;
			Write(structure, file, out sortedSites);
		}

		/// <summary>
		/// Writes VASPs POSCAR file. Proper sorting of the structure is handled automatically. The sorted structure is returned.
		/// </summary>
		public static void Write(Structure structure, FileInfo file, out IEnumerable<Site> sortedSites)
		{
			if (structure == null || structure.Sites.Count == 0)
			{
				throw new ArgumentException("Structure not defined.");
			}

			// List of atoms in the same order as to be specified in POSCAR and POTCAR.
			List<Atom> atoms = new List<Atom>();
			Dictionary<Atom, List<Site>> sites = new Dictionary<Atom, List<Site>>();
			List<int> order = new List<int>();

			foreach (Site site in structure.Sites)
			{
				Atom atom = site.Atom;

				if (atom is SpinPolarizedAtom)
				{
					atom = ((SpinPolarizedAtom)atom).BaseAtom;
				}

				if (!atoms.Contains(atom))
				{
					atoms.Add(atom);
					sites.Add(atom, new List<Site>());
				}

				sites[atom].Add(site);
			}

			sortedSites = new List<Site>();
			foreach (Atom atom in atoms)
			{
				((List<Site>)sortedSites).AddRange(sites[atom]);
			}

			// http://cms.mpi.univie.ac.at/vasp/guide/node59.html

			using (StreamWriter writer = new StreamWriter(file.FullName))
			{
				Lattice lattice = structure.Lattice;

				// Write chemical formula as VASP title in first line.
				writer.WriteLine(structure.SpinUnpolarize().Formula);

				// Transform lattice vectors as VASP only supports Cartesian coordinate systems.
				writer.WriteLine("1");
				foreach (SpaceVector p in new SpaceVector[] { lattice.A.PhysicalCoordinates, lattice.B.PhysicalCoordinates, lattice.C.PhysicalCoordinates })
				{
					writer.WriteLine(string.Format(CultureInfo.InvariantCulture, "{0:r} {1:r} {2:r}", p.X, p.Y, p.Z));
				}

				// Write symbols (this isn't used by VASP itself; VASP uses POTCAR).
				writer.WriteLine(string.Join(" ", atoms.ConvertAll<string>(atom => atom.Symbol)));

				// Count number of each different atom.
				writer.WriteLine(string.Join(" ", atoms.ConvertAll<string>(atom => sites[atom].Count.ToString())));

				// List of relative atomic positions in the same order as potentials in POTCAR.
				writer.WriteLine("Direct");
				foreach (Atom atom in atoms)
				{
					foreach (Site site in sites[atom])
					{
						// POSCAR positions should be relative to the lattice, not just the basis vectors.
						SpaceVector p = lattice.Decompose(site.Position).Coordinates;
						writer.WriteLine(string.Format(CultureInfo.InvariantCulture, "{0:r} {1:r} {2:r}", p.X, p.Y, p.Z));

						order.Add(structure.Sites.IndexOf(site));
					}
				}

				/*writer.WriteLine();
				writer.WriteLine("! Original structure");
				using (StringReader reader = new StringReader(new Structure(structure.Lattice, sortedSites).Write()))
				{
					string s;
					while ((s = reader.ReadLine()) != null)
					{
						writer.Write("! ");
						writer.WriteLine(s);
					}
					reader.Close();
				}
				writer.WriteLine("! Original order");
				writer.Write("! ");
				writer.WriteLine(string.Join(" ", order.ConvertAll<string>(index => index.ToString())));
				writer.WriteLine("! End of original structure");*/

				writer.Close();
			}
		}

		/// <summary>
		/// Reads VASPs POSCAR file.
		/// </summary>
		public static Structure Read(FileInfo file)
		{
			using (StreamReader reader = new StreamReader(file.FullName))
			{
				return Read(reader);
			}
		}

		/// <summary>
		/// Reads VASPs POSCAR file.
		/// </summary>
		public static Structure Read(TextReader reader)
		{
			return InnerRead(reader, null);
		}

		/// <summary>
		/// Reads VASPs POSCAR file.
		/// </summary>
		public static Structure Read(TextReader reader, IEnumerable<Atom> atoms)
		{
			if (atoms == null)
			{
				return Read(reader);
			}

			return InnerRead(reader, new List<Atom>(atoms));
		}

		private static Structure InnerRead(TextReader reader, List<Atom> atoms)
		{
			// Ignore first line.
			string line0 = reader.ReadLine();

			double scaling = double.Parse(reader.ReadLine(), CultureInfo.InvariantCulture);
			Lattice lattice = Lattice.Parse(AtomicEnvironment.ParseSpaceVector(reader.ReadLine()) * scaling, AtomicEnvironment.ParseSpaceVector(reader.ReadLine()) * scaling, AtomicEnvironment.ParseSpaceVector(reader.ReadLine()) * scaling);

			string line1 = reader.ReadLine().Trim();
			string line2;
			if (char.IsLetter(line1[0]))
			{
				// Everything is as expected. Read next line.
				line2 = reader.ReadLine();
			}
			else
			{
				// Try to extract atoms from the comment line instead. The is the output from Phonopy.
				line2 = line1;
				line1 = line0;
			}

			if (atoms == null)
			{
				// Read atomic symbols if not already supplied.
				atoms = new List<Atom>();
				foreach (string s in line1.Split(' '))
				{
					if (s != "")
					{
						if (!char.IsLetter(s[0]))
						{
							throw new Exception("Atomic symbols most be specified in the POSCAR file.");
						}

						atoms.Add(Atom.Parse(s));
					}
				}
			}

			List<int> counts = new List<int>();
			foreach (string s in line2.Split(' '))
			{
				if (s != "")
				{
					counts.Add(int.Parse(s));
				}
			}
		
			if (!reader.ReadLine().StartsWith("D"))
			{
				throw new Exception("Expected direct representation.");
			}

			List<Site> sites = new List<Site>();
			while (atoms.Count > 0)
			{
				Atom atom = atoms[0];
				counts[0] = counts[0] - 1;
				if (counts[0] == 0)
				{
					atoms.RemoveAt(0);
					counts.RemoveAt(0);
				}

				List<string> a = new List<string>();
				foreach (string s in reader.ReadLine().Split(' '))
				{
					if (s != "")
					{
						a.Add(s);
					}
				}

				Position position = new Position(lattice, AtomicEnvironment.ParseSpaceVector(string.Join(" ", a[0], a[1], a[2])));
				sites.Add(new Site(atom, position));
			}

			return new Structure(lattice, sites);
		}
	}
}
